from flask import Flask, render_template, jsonify, request
from flask_cors import CORS
import pickle
import networkx as nx
import json

app = Flask(__name__)
CORS(app)

# Load graph data
print("Loading graph data...")
with open("business-payout.pkl", "rb") as f:
    G = pickle.load(f)
print(f"Graph loaded: {G.number_of_nodes()} nodes, {G.number_of_edges()} edges")


def get_node_info(node_id):
    """Get detailed information about a node"""
    if node_id not in G:
        return None

    info = {
        "id": node_id,
        "in_degree": G.in_degree(node_id),
        "out_degree": G.out_degree(node_id),
        "total_degree": G.degree(node_id),
        "predecessors": list(G.predecessors(node_id)),
        "successors": list(G.successors(node_id)),
    }
    return info


def get_subgraph(center_node, max_degree=None, depth=1):
    """Get subgraph around a center node"""
    if center_node not in G:
        return None

    # Start with center node
    nodes = {center_node}

    # Add nodes within depth
    current_layer = {center_node}
    for _ in range(depth):
        next_layer = set()
        for node in current_layer:
            # Add predecessors and successors
            predecessors = list(G.predecessors(node))
            successors = list(G.successors(node))

            # Filter by degree if specified
            if max_degree:
                predecessors = [n for n in predecessors if G.degree(n) <= max_degree]
                successors = [n for n in successors if G.degree(n) <= max_degree]

            next_layer.update(predecessors[:20])  # Limit to prevent too many nodes
            next_layer.update(successors[:20])

        nodes.update(next_layer)
        current_layer = next_layer

        # Limit total nodes
        if len(nodes) > 200:
            break

    # Create subgraph
    subgraph = G.subgraph(nodes)

    # Convert to format suitable for visualization
    nodes_data = []
    for node in subgraph.nodes():
        node_attrs = G.nodes[node]
        nodes_data.append(
            {
                "id": node,
                "label": node.split("/")[-1] if "/" in node else node,
                "degree": G.degree(node),
                "in_degree": G.in_degree(node),
                "out_degree": G.out_degree(node),
                "group": 1 if node == center_node else 2,
                "type": node_attrs.get("type", "unknown"),
            }
        )

    edges_data = []
    for u, v, data in subgraph.edges(data=True):
        edges_data.append(
            {"source": u, "target": v, "type": data.get("type", "unknown")}
        )

    return {"nodes": nodes_data, "edges": edges_data, "center": center_node}


@app.route("/")
def index():
    return render_template("index.html")


@app.route("/api/search", methods=["GET"])
def search_nodes():
    """Search nodes by keyword"""
    query = request.args.get("q", "").lower()
    if not query:
        return jsonify([])

    # Search for nodes containing the query
    matching_nodes = []
    for node in G.nodes():
        if query in node.lower():
            matching_nodes.append(
                {
                    "id": node,
                    "label": node.split("/")[-1] if "/" in node else node,
                    "degree": G.degree(node),
                    "in_degree": G.in_degree(node),
                    "out_degree": G.out_degree(node),
                }
            )

    # Sort by degree and limit results
    matching_nodes.sort(key=lambda x: x["degree"], reverse=True)
    return jsonify(matching_nodes[:50])


@app.route("/api/node/<path:node_id>")
def get_node(node_id):
    """Get node information"""
    info = get_node_info(node_id)
    if info:
        return jsonify(info)
    return jsonify({"error": "Node not found"}), 404


@app.route("/api/graph/<path:node_id>")
def get_graph(node_id):
    """Get subgraph for visualization"""
    max_degree = request.args.get("max_degree", type=int)
    depth = request.args.get("depth", 1, type=int)

    graph_data = get_subgraph(node_id, max_degree, depth)
    if graph_data:
        return jsonify(graph_data)
    return jsonify({"error": "Node not found"}), 404


@app.route("/api/expand/<path:node_id>")
def expand_node(node_id):
    """Expand a node to show its connections"""
    if node_id not in G:
        return jsonify({"error": "Node not found"}), 404

    # Get direct connections
    predecessors = list(G.predecessors(node_id))[:10]
    successors = list(G.successors(node_id))[:10]

    nodes_data = []
    edges_data = []

    # Add the connections as nodes
    all_nodes = set(predecessors + successors + [node_id])
    for node in all_nodes:
        node_attrs = G.nodes[node]
        nodes_data.append(
            {
                "id": node,
                "label": node.split("/")[-1] if "/" in node else node,
                "degree": G.degree(node),
                "in_degree": G.in_degree(node),
                "out_degree": G.out_degree(node),
                "group": 1 if node == node_id else 2,
                "type": node_attrs.get("type", "unknown"),
            }
        )

    # Add edges
    for pred in predecessors:
        edges_data.append({"source": pred, "target": node_id, "type": "dependency"})

    for succ in successors:
        edges_data.append({"source": node_id, "target": succ, "type": "dependency"})

    return jsonify({"nodes": nodes_data, "edges": edges_data})


@app.route("/api/stats")
def get_stats():
    """Get graph statistics"""
    degrees = dict(G.degree())
    sorted_nodes = sorted(degrees.items(), key=lambda x: x[1], reverse=True)

    stats = {
        "total_nodes": G.number_of_nodes(),
        "total_edges": G.number_of_edges(),
        "avg_degree": sum(degrees.values()) / len(degrees) if degrees else 0,
        "top_nodes": [
            {
                "id": node,
                "label": node.split("/")[-1] if "/" in node else node,
                "degree": degree,
                "in_degree": G.in_degree(node),
                "out_degree": G.out_degree(node),
            }
            for node, degree in sorted_nodes[:20]
        ],
    }
    return jsonify(stats)


if __name__ == "__main__":
    app.run(debug=True, host="0.0.0.0", port=5555)